from torch import nn
import torch
from torch.nn import functional as F
class MySqu(nn.Module):
    def __init__(self,*args):
        super().__init__()
        for block in args:
            self.modules[block] = block
    
    def forward(self,X):
        for block in self._modules.values():
            X = block(X)
        return X

class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.hidden = nn.Linear(20,256)
        self.output = nn.Linear(256,10)
    
    def forward(self,X):
        return self.output(F.relu(self.hidden(X)))

net = MLP()
X = torch.randn(size=(2,20))
Y = net(X)
print(Y)

torch.save(net.state_dict(),'mlp.params')
clone = MLP()
clone.load_state_dict(torch.load('mlp.params'))
clone.eval()
Y_clone = clone(X)
print(Y_clone==Y)

